#!/usr/bin/python

"""
Python program for interacting with the VirusTotal.com public API.  You must have a
virustotal.com account and a public API key to use this program.

post_multipart(), encode_mutlipart_formdata(), and some other useful lines of code
provided by VirusTotal.com.  See http://www.virustotal.com/advanced.html

Copyright 2010 Paul Edelman
GPLv3

Author:   Paul Edelman
Modified: 16 SEP 2010
"""

import sys
import getopt
import json
import urllib
import urllib2
import time
import httplib
import mimetypes

"""
Global Variables
"""
# VT does not allow you to query more often than every 15 seconds.
sleepTime = 15
# Make sure to set this to your VT public API key.
key = ""

"""
Postfile Code
"""

def post_multipart(host, selector, fields, files):
    """
    Post fields and files to an http host as multipart/form-data.
    fields is a sequence of (name, value) elements for regular form fields.
    files is a sequence of (name, filename, value) elements for data to be uploaded as files
    Return the server's response page.
    """
    content_type, body = encode_multipart_formdata(fields, files)
    h = httplib.HTTPS(host)
    h.putrequest('POST', selector)
    h.putheader('content-type', content_type)
    h.putheader('content-length', str(len(body)))
    h.endheaders()
    h.send(body)
    errcode, errmsg, headers = h.getreply()
    return h.file.read()

def encode_multipart_formdata(fields, files):
    """
    fields is a sequence of (name, value) elements for regular form fields.
    files is a sequence of (name, filename, value) elements for data to be uploaded as files
    Return (content_type, body) ready for httplib.HTTP instance
    """
    BOUNDARY = '----------ThIs_Is_tHe_bouNdaRY_$'
    CRLF = '\r\n'
    L = []
    for (key, value) in fields:
        L.append('--' + BOUNDARY)
        L.append('Content-Disposition: form-data; name="%s"' % key)
        L.append('')
        L.append(value)
    for (key, filename, value) in files:
        L.append('--' + BOUNDARY)
        L.append('Content-Disposition: form-data; name="%s"; filename="%s"' % (key, filename))
        L.append('Content-Type: %s' % get_content_type(filename))
        L.append('')
        L.append(value)
    L.append('--' + BOUNDARY + '--')
    L.append('')
    body = CRLF.join(L)
    content_type = 'multipart/form-data; boundary=%s' % BOUNDARY
    return content_type, body

def get_content_type(filename):
    return mimetypes.guess_type(filename)[0] or 'application/octet-stream'

"""
usage()
Prints a program usage message
"""
def usage():
    print "virustotal.py --mode <mode> --type <type> --output <outputFile> <[inputFile0,inputFile1,...,inputFileN]>"
    print ""
    print "Arguments:"
    print ""
    print "  -m --mode          Operating mode [report, scan]"
    print "  -t --type          Mode type [file, url]"
    print "  -o --output        Path of the file to write output to"
    print "  --key-file         File containing the API key, defaults to api_key.txt in working directory"
    print "  -h --help          Prints this help message."
    print ""
    print "Examples:"
    print "  virustotal.py --mode report --type file --output <outputFile> <resourceIDList.txt>"
    print "    Takes as input a newline separated list of resource IDs. Resource IDs"
    print "    include MD5 or the scan IDs produced from submitting several files to"
    print "    be scanned."
    print ""
    print "  virustotal.py --mode report --type url --output <outputFile> <urlList.txt>"
    print "    Takes as input a newline separated list of URLs. Returns a single list"
    print "    of CSV reports, one report per URL per line."
    print ""
    print "  virustotal.py --mode scan --type file --output <outputFile> <[file0,file1,file2,...,fileN]>"
    print "  virustotal.py --mode scan --type url --output <outputFile> <urlList.txt>"

"""
help()
Prints the standard help message.
"""
def help():
   print "This program is a Python client which makes use of the virustotal.com public API."
   print "In order to use this program you will need to acquire an API key from virustotal.com"
   print "which requires a user account."
   print ""
   usage()

"""
get_request_string()
Creates a request string based on a resource ID.  The resource ID could be an
MD5 for some other internal value.
"""
def get_request_string(api_url, resource):
    parameters = {"resource": resource,
                  "key": key}
    data = urllib.urlencode(parameters)
    request = urllib2.Request(api_url, data)
    return request

"""
retrieve_scan_report()
Sends a scan report request for the given MD5 and returns the
properly formatted output.
"""
def retrieve_file_scan_report(inputFileName, outputFileName):
    api_url = "https://www.virustotal.com/api/get_file_report.json"

    # Read in resource IDs from the input file and process them one at a time
    try:
        f = open(inputFileName)
        fout = open(outputFileName, 'w')

        # Begin the resource report processing loop
        numRequests = 0
        for line in f:
            resourceID = line.strip()

            # Break on all empty lines
            if resourceID == "":
                break
            else:
                numRequests += 1

            print "Processing Resource: " + resourceID

            # Get the formatted output string for the report on the given MD5
            httpResponse = None

            # Create the HTTPRequest, and get the HTTPResponse
            try:
                httpResponse = urllib2.urlopen(get_request_string(api_url, resourceID))
            except Exception as e:
                errorMsg = resourceID + ",error," + str(e)
                print resourceID + ": " + str(e)
                return errorMsg

            # Process the HTTPResponse and get the JSON object
            jsonObject = httpResponse.read()
            report = json.loads(jsonObject).get("report")

            # By default the output_string is just the resourceID
            output_string = resourceID

            # Add report data to the output string if we have a report, otherwise
            # just give the standard report string defined by default.
            # TODO Some code replication here, consider consolidation.
            if report is not None:

                # Parse the report
                timeStamp = report[0]
                reportEntries = report[1]
                numEntries = len(reportEntries)

                # Determine number of detections
                numDetects = 0
                entryValues = dict.values(reportEntries)
                for v in entryValues:
                    if v != u'':
                        numDetects += 1

                # Create and return the output string
                output_string = resourceID + ",1," + str(numDetects) + "," + str(numEntries) + "," + timeStamp
                for k,v in reportEntries.iteritems():
                    k = k.encode("ascii")
                    v = v.encode("ascii")
                    if v == "":
                        v = "-"
                    output_string += "," + k + "," + v
            else:
                output_string = output_string + ",0"

            # Write output_string to output file
            fout.write(output_string + "\n")
            # This increases file I/O, but means we don't have to wait to start seeing
            # results. Should consider the performance implications of this.
            fout.flush()

            # Sleep for a bit so as not to overload virustotal.com
            time.sleep(sleepTime)
        # Done with the resource report processing loop

        print "Processed " + str(numRequests) + " requests."

        # Close both the input and output files
        fout.write('\n')
        fout.close()
        f.close()
    except IOError:
        print "Error: Unable to open the input file specified"
        sys.exit(1)

"""
send_and_scan_file()
Uploads a file to virustotal.com for scanning.
"""
def send_and_scan_files(filenames, outputFileName):
    api_url = "/api/scan_file.json"
    host = "www.virustotal.com"

    fields = [("key", key)]

    try:
        # Write the scan IDs to the specified output file
        fdata = open(outputFileName, "w")
        flog = open(outputFileName + ".log", "w")

        # Begin file scan processing loop
        numFiles = len(filenames)
        count = 0
        for filename in filenames:
            count += 1
            print "Sending: " + filename
            file_to_send = open(filename, "rb").read()
            files = [("file", filename, file_to_send)]
            jsonObject = post_multipart(host, api_url, fields, files)
            report = json.loads(jsonObject)
            resultCode = report['result']
            if resultCode == True:
                print "Successful"
                fdata.write(report['scan_id'] + "\n")
                flog.write(filename + "," + str(resultCode) + "," + report['scan_id'] + "\n")
            else:
                print "Unsuccessful"
                flog.write(filename + "," + str(resultCode) + "\n")

            # Sleep between requests
            if count < numFiles:
                time.sleep(sleepTime)

    except Exception as e:
        print "Error: ", e

    print "Processed " + str(count) + " files."

    # Close output files
    fdata.close()
    flog.close()

"""
retrieve_url_scan_report()
Sends a request to get the scan report for the given URL.
"""
def retrieve_url_scan_report(inputFileName, outputFileName):
    api_url = "https://www.virustotal.com/api/get_url_report.json"

    try:
        # Open input file
        fin = open(inputFileName)
        fout = open(outputFileName, "w")

        # Begin loop handling all URLs listed in the input file
        numURLs = 0
        for url in fin:
            url = url.strip()

            # Break on all empty lines
            if url == "":
                break
            else:
                numURLs += 1

            # Begin processing this URL
            print "Processing URL: " + url
            response = urllib2.urlopen(get_request_string(api_url, url))
            jsonObj = response.read()

            # Format the report
            report = json.loads(jsonObj).get("report")
            output_string = url

            # TODO Some code replication here, consider consolidation.
            if report is not None:

                # Parse the report
                timeStamp = report[0]
                reportEntries = report[1]
                numEntries = len(reportEntries)

                # Determine number of detections
                numDetects = 0
                entryValues = dict.values(reportEntries)
                for v in entryValues:
                    if v != u'Clean site':
                        numDetects += 1

                # Create formatted report string
                output_string = output_string + ",1," + str(numDetects) + "," + str(numEntries) + "," + timeStamp
                for k,v in reportEntries.iteritems():
                    k = k.encode("ascii")
                    v = v.encode("ascii")
                    output_string += "," + k + "," + v
            else:
                output_string = output_string + ",0"

            # Write formatted report to output file
            fout.write(output_string + "\n")
            fout.flush()

            # Sleep between requests
            time.sleep(sleepTime)

        # Give some summary information
        print "Processed " + str(numURLs) + " URLs"

    except Exception as e:
        print "Error: ", e

    # Close files
    fout.close()
    fin.close()

"""
send_and_scan_url()
Sends a request to scan the given URL and returns the report result, but first
checks to see if the URL already has a scan report.
"""
def send_and_scan_url(inputFileName, outputFileName):
    api_url = "https://www.virustotal.com/api/scan_url.json"

    # TODO Code replication with retrieve_url_scan_report(), consider refactoring
    try:
        fin = open(inputFileName)
        fdata = open(outputFileName, "w")
        flog = open(outputFileName + ".log", "w")

        numURLs = 0
        for url in fin:
            url = url.strip()

            # Break on all empty lines
            if url == "":
                break
            else:
                numURLs += 1

            # Process this URL
            print "Processing URL: " + url
            response = urllib2.urlopen(get_request_string(api_url, url))
            jsonObj = response.read()
            print jsonObj
            report = json.loads(jsonObj)

            resultCode = report['result']
            if resultCode == True:
                print "Successful"
                fdata.write(report['scan_id'] + "\n")
                flog.write(url + "," + str(resultCode) + "," + report['scan_id'] + "\n")
            else:
                print "Unsuccessful"
                flog.write(url + "," + str(resultCode) + "\n")

    except Exception as e:
        print "Error: ", e

    # Close files
    flog.close()
    fdata.close()
    fin.close()

##################################################
# BEGIN MAIN EXECUTION

# Check for valid number of arguments
#if len(sys.argv) < 5:
#    print "Error: Invalid number of arguments"
#    usage()
#    sys.exit(1)

keyFileName = "api_key.txt"
mode = ""
type = ""
outputFileName = ""

# Process arguments
opts, args = getopt.getopt(sys.argv[1:], "m:t:o:h", ["key-file=", "mode=", "type=", "output=", "help"])

for o, a in opts:
    if o == 'key-file':
        keyFileName = a
    elif o in ('-m', '--mode'):
        mode = a
    elif o in ('-t', '--type'):
        type = a
    elif o in ('-o', '--output'):
        outputFileName = a
    elif o in ('-h', '--help'):
        help()
        sys.exit(1)
    else:
        pass

# Read in VT public API key from file
try:
    f = open(keyFileName, 'r')
    key = f.read().strip()
    print "Using API Key: " + key
    f.close()
except IOError:
    print "Error: Unable to open the api_key.txt file"
    sys.exit(1)

# Check for valid AT API key
if key == "":
    print "Error: Must provide a valid API key in the file api_key.txt in this directory."
    usage()
    sys.exit(1)

# Check for valid sleep time
if sleepTime < 15:
    print "Error: Must sleep for at least 15 seconds between queries, please increase the sleep time"
    usage()
    sys.exit(1)

# Process mode and type arguments
if mode == "report":
    # Retrieving a report
    if type == "file":
        inputFileName = args[0]
        print "Operation: File Scan Report"
        print "Reading input from: " + inputFileName
        print "Storing results in: " + outputFileName
        retrieve_file_scan_report(inputFileName, outputFileName)
    elif type == "url":
        print "Operation: URL Scan Report"
        inputFileName = args[0]
        retrieve_url_scan_report(inputFileName, outputFileName)
    else:
        print "Error: Invalid <type> specified"
        usage()
        sys.exit(1)
elif mode == "scan":
    if type == "file":
        print "Operation: Send and Scan File(s)"
        # We could have one to many files listed on the command line
        send_and_scan_files(args, outputFileName)
    elif type == "url":
        print "Operation: Send and Scan URL(s)"
        inputFileName = args[0]
        send_and_scan_url(inputFileName, outputFileName)
    else:
        print "Error: Invalid <type> specified"
        usage()
        sys.exit(1)
else:
    print "Error: Invalid <mode> specified"
    usage()
    sys.exit(1)

# END OF MAIN EXECUTION
##################################################

